import re
from enum import Enum, auto

import pandas as pd
from PIL import Image as PILImage

from ai_translator_v2.utils import LOG


class ContentType(Enum):
    TEXT = auto()
    TABLE = auto()
    IMAGE = auto()


class Content(object):
    def __init__(self, content_type, original, translation=None):
        self.content_type = content_type
        self.original = original
        self.translation = translation
        self.status = False

    def set_translation(self, translation, status):
        if not self.check_translation_type(translation):
            raise ValueError(f"Invalid translation type. Expected {self.content_type}, but got {type(translation)}")
        LOG.debug(f"[transaction]\n{translation}")
        self.translation = translation
        self.status = status

    def check_translation_type(self, translation):
        if self.content_type == ContentType.TEXT and isinstance(translation, str):
            return True
        if self.content_type == ContentType.TABLE and isinstance(translation, list):
            return True
        if self.content_type == ContentType.IMAGE and isinstance(translation, PILImage.Image):
            return True
        return False

    def __str__(self):
        return self.original


class TableContent(Content):
    def __init__(self, data, translation=None):
        df = pd.DataFrame(data)
        if len(data) != len(df) or len(data[0]) != len(df.columns):
            raise ValueError(f"The number of rows and columns in the extracted table do not match")
        super().__init__(ContentType.TABLE, df)

    def set_translation(self, translation, status):
        """
        设置翻译结果到类属性中。

        尝试解析提供的翻译字符串为一个DataFrame，并更新类的status和translation属性。
        如果解析过程中出现错误，将记录错误并重置translation为None，status为False。

        参数:
        - translation: str，一个包含翻译数据的字符串，预期格式为一个由逗号或中文逗号分隔的列表，
                       列名在第一个方括号之后，数据行在第一个方括号之后且每个数据行由方括号包围。
        - status: bool，表示翻译结果是否有效或可用。

        返回值: 无
        """
        try:
            # 检查translation参数是否为字符串类型
            if not isinstance(translation, str):
                raise ValueError(f"Invalid translation type. Expected str, but got {type(translation)}")

            # 记录翻译字符串的调试信息
            LOG.debug(f"[transaction]\n{translation}")

            # 解析翻译字符串中的头部信息（列名）和数据行
            header = re.split('[,，]', translation.split(']')[0].strip()[1:])
            data_rows = translation.strip().split(']')[1:]
            data_rows = [re.split('[,，]', row.strip()[1:-1]) for row in data_rows if row]

            # 将解析出的数据行转换为DataFrame
            translated_df = pd.DataFrame(data_rows, columns=header)

            # 记录转换后的DataFrame的调试信息
            LOG.debug(f"f[translated_df]\n{translated_df}")

            # 更新类属性
            self.translation = translated_df
            self.status = status
        except Exception as e:
            # 如果在处理过程中发生异常，记录错误并重置translation和status
            LOG.error(f"An error occurred during table translation: {e}")
            self.translation = None
            self.status = False

    def __str__(self):
        return self.original.to_string(header=False, index=False)

    def iter_items(self, translated=False):
        """
        遍历数据集中的项。

        此函数提供了一个迭代器，用于逐行遍历数据集（可以选择是原始数据集还是翻译后的数据集），并逐个返回每一行的索引、列索引和对应的值。

        参数:
        translated (bool): 指定是否使用翻译后的数据集。默认为False，表示使用原始数据集。

        返回值:
        generator: 生成器对象，每次迭代返回一个三元组 (row_idx, column_idx, column)，分别表示行索引、列索引和该位置上的值。
        """
        # 根据 translated 参数选择遍历的目标数据集
        target_df = self.translation if translated else self.original
        # 遍历目标数据集的每一行
        for row_idx, row in target_df.iterrows():
            # 遍历当前行的每一列
            for column_idx, column in enumerate(row):
                # 生成并yield每一行的索引、列索引和对应的值
                yield row_idx, column_idx, column

    def update_item(self, row_idx, column_idx, new_value, translated=False):
        target_df = self.translation if translated else self.original
        target_df.at[row_idx, column_idx] = new_value

    def get_original_as_str(self):
        return self.original.to_string(header=False, index=False)
